import random
import os

VID_CLASSES_name =(  # always index 0
    'airplane', #1 airplane
    'antelope', #2 antelope
    'bear', #3 bear
    'bicycle', #4 bicycle
    'bird', #5 bird
    'bus', #6 bus
    'car', #7 car
    'cattle', #8 cattle
    'dog', #9 dog
    'domestic_cat', #10 domestic_cat
    'elephant', #11 elephant
    'fox', #12 fox
    'giant_panda', #13 giant_panda
    'hamster', #14 hamster
    'horse', #15 horse
    'lion', #16 lion
    'lizard', #17 lizard
    'monkey', #18 monkey
    'motorcycle', #19 motorcycle
    'rabbit', #20 rabbit
    'red_panda', #21 red_panda
    'sheep', #22 sheep
    'snake', #23 snake
    'squirrel', #24 squirrel
    'tiger', #25 tiger
    'train', #26 train
    'turtle', #27 turtle
    'watercraft', #28 watercraft
    'whale', #29 whale
    'zebra', #30 zebra
)

DET_file = '/home/sean/data/ILSVRC/ImageSets/DET/train_pos'
DET_num = [2,4,20,24,26,33,37,39,58,59,64,70,74,84,92,103,105,113,114,141,144,155,159,166,182,185,188,197,198,200]

VID_file = '/home/sean/data/ILSVRC/ImageSets/VID/train_video_remove_no_object.txt'
write_file = '/home/sean/data/ILSVRC/ImageSets/VID/train_VID_DET.txt'

DET_hat = '../../DET/ILSVRC2014_DET_train/'

num_each_class_DET = 2000
num_each_class_VID = 10
DET_total_num = 0
with open(write_file, 'w') as wf:
    for i in DET_num:
        with open(DET_file+'_'+ str(i)+'.txt', 'r') as f:
            lines = f.readlines()
            if len(lines)<=num_each_class_DET:
                for line in lines:
                    wf.write(DET_hat+line)
                print(VID_CLASSES_name[DET_num.index(i)], len(lines))
                DET_total_num += len(lines)
            else:
                skip = int(len(lines) / num_each_class_DET)+1
                select_list = lines[0:len(lines):skip]
                for line in select_list:
                    wf.write(DET_hat+line)
                DET_total_num += len(select_list)
                print(VID_CLASSES_name[DET_num.index(i)], len(select_list))
    print('DET', DET_total_num)
    for line in open(VID_file, 'r'):
        video_name, video_length = line.split(' ')
        video_len = int(video_length)
        if video_len <= num_each_class_VID:
            for id in range(video_len):
                wf.write(os.path.join(video_name, str(id).zfill(6))+'\n')
        else:
            skip = int(video_len / num_each_class_VID)
            uniform_list = list(range(0, video_len, skip))
            cast_list = random.sample(range(len(uniform_list)), len(uniform_list) - num_each_class_VID)
            select_list = [x for x in uniform_list if uniform_list.index(x) not in cast_list]
            for id in select_list:
                wf.write(os.path.join(video_name, str(id).zfill(6)) + '\n')
            print(select_list, len(select_list))

